from arm import *
import util_3d as u3d
from spatialmath.base import *
from d435 import D435
import halcon as ha
import cv2


def gen_pose_to_record_end(org_posture_rokea_type):
    pose_array = []
    base_t_end = u3d.rokae_pose_to_T(org_posture_rokea_type)
    # end_t_calib = u3d.xyzrpy_to_T(-0.175, 0.11, 0.08, 0.0, 0.0, 0.0)
    end_t_calib = u3d.xyzrpy_to_T(-0.0, -0.0, 0.20, 0.0, 0.0, 0.0)
    calib_t_end = trinv(end_t_calib)
    calib_aim_t_end_aim = calib_t_end
    for i in range(10):
        delta_calib = u3d.xyzrpy_to_T(0.0, 0.0, 0.0, 0.0, 0, -10+i*2)
        calib_t_calib_aim = delta_calib
        base_t_end_aim = base_t_end @ delta_calib
        base_t_end_aim_rokae = u3d.T_to_rokae_pose(base_t_end_aim)
        pose_array.append(base_t_end_aim_rokae)

    for i in range(10):
        delta_calib = u3d.xyzrpy_to_T(0.0, 0.0, 0.0, 0.0, -10+i*2, 0)
        calib_t_calib_aim = delta_calib
        base_t_end_aim = base_t_end @ delta_calib
        base_t_end_aim_rokae = u3d.T_to_rokae_pose(base_t_end_aim)
        pose_array.append(base_t_end_aim_rokae)


    for i in range(10):
        delta_calib = u3d.xyzrpy_to_T(0.0, 0.0, 0.0, -10+i*2, 0, 0)
        calib_t_calib_aim = delta_calib
        base_t_end_aim = base_t_end @ delta_calib
        base_t_end_aim_rokae = u3d.T_to_rokae_pose(base_t_end_aim)
        pose_array.append(base_t_end_aim_rokae)

    for i in range(5):
        delta_calib = u3d.xyzrpy_to_T(-0.025+0.01*i, 0.0, 0.0, 0.0, 0, 0)
        calib_t_calib_aim = delta_calib
        base_t_end_aim = base_t_end @ delta_calib
        base_t_end_aim_rokae = u3d.T_to_rokae_pose(base_t_end_aim)
        pose_array.append(base_t_end_aim_rokae)


    for i in range(5):
        delta_calib = u3d.xyzrpy_to_T(0, -0.025+0.01*i, 0.0, 0.0, 0, 0)
        calib_t_calib_aim = delta_calib
        base_t_end_aim = base_t_end @ delta_calib
        base_t_end_aim_rokae = u3d.T_to_rokae_pose(base_t_end_aim)
        pose_array.append(base_t_end_aim_rokae)

    for i in range(5):
        delta_calib = u3d.xyzrpy_to_T(0, 0.0, -0.025+0.01*i, 0.0, 0, 0)
        calib_t_calib_aim = delta_calib
        base_t_end_aim = base_t_end @ delta_calib
        base_t_end_aim_rokae = u3d.T_to_rokae_pose(base_t_end_aim)
        pose_array.append(base_t_end_aim_rokae)

    return pose_array
dir_name = './data'
if __name__ == "__main__":
    arm = Rokae()
    arm.init()
    camera_d435 = D435()
    print(arm.getPose())
    p2 = arm.getPose()
    arm.setPose2(lambda st: True, 50, 50, 50, p2)
    pose_list = gen_pose_to_record_end(p2)
    for i in range(30):
        alpha = 57.295779513
        time.sleep(1)
        pose_file_name = f"{i}.dat"
        image_file_name = f"{i}.png"
        pose_file_name = dir_name + '/' + pose_file_name
        image_file_name = dir_name + '/' + image_file_name
        arm.setPose(lambda st:True, pose_list[i])
        posture = arm.getPose()
        camera_d435.update()
        img,img_depth = camera_d435.get_images()
        cv2.imwrite(image_file_name, img)
        cv2.imshow("tmp", img)
        cv2.waitKey(10)
        # ha.write_image(img, 'png', 0, image_file_name)
        print(f"Record data  {i}")
        with open(pose_file_name, 'w') as f:
            text1 = f"r {alpha * posture[3]} {alpha * posture[4]} {alpha * posture[5]}\n"
            text2 = f"t {posture[0]} {posture[1]} {posture[2]}\n"
            f.write('f 2\n')
            f.write(text1)
            f.write(text2)
    print(f"Record data end")
    arm.setPose2(lambda st: True, 50, 50, 50, p2)
    time.sleep(2)
    arm.close()
    pass
